from captcha.image import ImageCaptcha
import numpy as np
import matplotlib.pyplot as plt
from PIL import Image
import random
import tensorflow as tf
import configparser
import os,sys

number = ['0', '1', '2', '3', '4', '5', '6', '7', '8', '9']
alphabet = ['a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n', 'o', 'p', 'q', 'r', 's', 't', 'u',
            'v', 'w', 'x', 'y', 'z']
Alphabet = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P', 'Q', 'R', 'S', 'T', 'U',
            'V', 'W', 'X', 'Y', 'Z']

char_set = number  # + alphabet + Alphabet

# 图片高
IMAGE_HEIGHT = 60
# 图片宽
IMAGE_WIDTH = 160
# 验证码长度
MAX_CAPTCHA = 4
# 验证码选择空间
CHAR_SET_LEN = len(char_set)
# 提前定义变量空间
X = tf.placeholder(tf.float32, [None, IMAGE_HEIGHT * IMAGE_WIDTH])
Y = tf.placeholder(tf.float32, [None, MAX_CAPTCHA * CHAR_SET_LEN])
keep_prob = tf.placeholder(tf.float32)  # 节点保留率

def ini_read(filename, node, name):
    config = configparser.ConfigParser()
    config.readfp(open(filename))
    text = config.get(node, name)
    return text

def ini_write(filename, node, name):
    config = configparser.ConfigParser()
# set a number of parameters
    config.add_section("book")
    config.set("book", "title", "the python standard library")
    config.set("book", "author", "fredrik lundh")
    config.add_section("ematter")
    config.set("ematter", "pages", 250)
# write to file
    config.write(open('1.ini', "w"))
    return 0

def ini_change(filename, nade, name, value):
    config = configparser.ConfigParser()
    config.read(filename)
    config.set(nade, name, value)
    config.write(open(filename, "r+"))  # 可以把r+改成其他方式，看看结果:)
    return 0

# 生成n位验证码字符 这里n=4
def random_captcha_text(char_set=char_set, captcha_size=4):
    captcha_text = []
    for i in range(captcha_size):
        c = random.choice(char_set)
        captcha_text.append(c)
    return captcha_text

# 使用ImageCaptcha库生成验证码
def gen_captcha_text_and_image():
    image = ImageCaptcha()
    captcha_text = random_captcha_text()
    captcha_text = ''.join(captcha_text)
    captcha = image.generate(captcha_text)
    captcha_image = Image.open(captcha)
    captcha_image = np.array(captcha_image)
    return captcha_text, captcha_image

# 彩色图转化为灰度图
def convert2gray(img):
    if len(img.shape) > 2:
        gray = np.mean(img, -1)
        # 上面的转法较快，正规转法如下
        # r, g, b = img[:,:,0], img[:,:,1], img[:,:,2]
        # gray = 0.2989 * r + 0.5870 * g + 0.1140 * b
        return gray
    else:
        return img

# 获取字符在 字符域中下标
def getPos(char_set=char_set, char=None):
    return char_set.index(char)

# 验证码字符转换为长向量
def text2vec(text):
    text_len = len(text)
    if text_len > MAX_CAPTCHA:
        raise ValueError('验证码最长4个字符')

    vector = np.zeros(MAX_CAPTCHA * CHAR_SET_LEN)

    def char2pos(c):
        if c == '_':
            k = 62
            return k
        k = ord(c)-48
        if k > 9:
            k = ord(c) - 55
            if k > 35:
                k = ord(c) - 61
                if k > 61:
                    raise ValueError('No Map')
        return k

    for i, c in enumerate(text):
        idx = i * CHAR_SET_LEN + char2pos(c)
        #idx = i * CHAR_SET_LEN + getPos(char=c)
        vector[idx] = 1
    return vector

# 获得1组验证码数据
def get_next_batch(batch_size=128):
    batch_x = np.zeros([batch_size, IMAGE_HEIGHT * IMAGE_WIDTH])
    batch_y = np.zeros([batch_size, MAX_CAPTCHA * CHAR_SET_LEN])

    def wrap_gen_captcha_text_and_image():
        while True:
            text, image = gen_captcha_text_and_image()
            if image.shape == (60, 160, 3):
                return text, image

    for i in range(batch_size):
        text, image = wrap_gen_captcha_text_and_image()
        image = convert2gray(image)
        batch_x[i, :] = image.flatten() / 255
        batch_y[i, :] = text2vec(text)
    return batch_x, batch_y

# 卷积层 附relu  max_pool drop操作
def conn_layer(w_alpha=0.01, b_alpha=0.1, _keep_prob=0.7, input=None, last_size=None, cur_size=None):
    w_c = tf.Variable(w_alpha * tf.random_normal([3, 3, last_size, cur_size]))
    b_c = tf.Variable(b_alpha * tf.random_normal([cur_size]))
    
    #卷积函数
    #1 input
    #2 卷积核 filter
    #3 步长  strides
    #4 padding：string类型的量，只能是"SAME","VALID
    #5 use_cudnn_on_gpu:bool类型，是否使用cudnn加速，默认为true
    conv = tf.nn.conv2d( input, w_c, strides=[1, 1, 1, 1], padding='SAME')
    
    #将偏差项 bias 加到 value 上面
    conv = tf.nn.bias_add(conv, b_c)
    #激活函数
    conv = tf.nn.relu(conv)
    #池化
    conv = tf.nn.max_pool(conv, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')
    #防止过拟合
    conv = tf.nn.dropout(conv, keep_prob=_keep_prob)
    return conv

# 对卷积层到全链接层的数据进行变换
def _get_conn_last_size(input):
    shape = input.get_shape().as_list()
    dim = 1
    for d in shape[1:]:
        dim *= d
    input = tf.reshape(input, [-1, dim])
    return input, dim

# 全链接层
def _fc_layer(w_alpha=0.01, b_alpha=0.1, input=None, last_size=None, cur_size=None):
    w_d = tf.Variable(w_alpha * tf.random_normal([last_size, cur_size]))
    b_d = tf.Variable(b_alpha * tf.random_normal([cur_size]))
    fc = tf.nn.bias_add(tf.matmul(input, w_d), b_d)
    return fc

# 构建前向传播网络
def crack_captcha_cnn(w_alpha=0.01, b_alpha=0.1):
    # TF-调整矩阵维度
    x = tf.reshape(X, shape=[-1, IMAGE_HEIGHT, IMAGE_WIDTH, 1])
    # 卷积层 附relu  max_pool drop操作
    conv1 = conn_layer(input=x, last_size=1, cur_size=32)
    conv2 = conn_layer(input=conv1, last_size=32, cur_size=64)
    conn3 = conn_layer(input=conv2, last_size=64, cur_size=64)

    input, dim = _get_conn_last_size(conn3)  # 对卷积层到全链接层的数据进行变换
    # 全链接层
    fc_layer1 = _fc_layer(input=input, last_size=dim, cur_size=1024)
    fc_layer1 = tf.nn.relu(fc_layer1)
    fc_layer1 = tf.nn.dropout(fc_layer1, keep_prob)

    fc_out = _fc_layer(input=fc_layer1, last_size=1024,
                       cur_size=MAX_CAPTCHA * CHAR_SET_LEN)
    return fc_out

# 反向传播
def back_propagation():
    output = crack_captcha_cnn()
    # 学习率
    back_loss = tf.reduce_mean(
        tf.nn.sigmoid_cross_entropy_with_logits(labels=Y, logits=output))
    back_optm = tf.train.AdamOptimizer(learning_rate=0.001).minimize(back_loss)
    predict = tf.reshape(output, [-1, MAX_CAPTCHA, CHAR_SET_LEN])
    max_idx_p = tf.arg_max(predict, 2)
    max_idx_l = tf.arg_max(tf.reshape(Y, [-1, MAX_CAPTCHA, CHAR_SET_LEN]), 2)
    back_accuracy = tf.reduce_mean(
        tf.cast(tf.equal(max_idx_p, max_idx_l), tf.float32))
    return back_loss, back_optm, back_accuracy

#加载现有模型 继续进行训练
def train_continue(step, acc):

    loss, optm, accuracy = back_propagation()

    saver = tf.train.Saver(max_to_keep=5)
    with tf.Session() as sess:
        if step == 0:
            sess.run(tf.global_variables_initializer())
        else:
            file_path = "models/crack_capcha.model-" + str(step)
            saver.restore(sess, file_path)
        # 36300 36300 0.9325 0.0147698

        acc1 = acc
        while 1:
            batch_x, batch_y = get_next_batch(64)
            _, loss_ = sess.run([optm, loss], feed_dict={
                                X: batch_x, Y: batch_y, keep_prob: 0.75})
            print(step, acc1, loss_)
            if step % 10 == 0:
                batch_x_test, batch_y_test = get_next_batch(100)
                acc1 = sess.run(accuracy, feed_dict={
                                X: batch_x_test, Y: batch_y_test, keep_prob: 1.})
                print(step, acc1, loss_)

                if acc1 > acc:
                    saver.save(sess, 'models/crack_capcha.model',
                               global_step=step)
                    ini_change(path+"/step.ini",
                               "value", "acc", str(acc1))
                    ini_change(path+"/step.ini",
                               "value", "step", str(step))
                    print(step, acc1, loss_, "保存")
                    acc = acc1
                if acc >= 0.99:
                    #saver.save(sess, 'models/crack_capcha.model', global_step=step)
                    break

            step += 1

# 测试训练模型
def crack_captcha(captcha_image, step):
    output = crack_captcha_cnn()

    saver = tf.train.Saver()
    with tf.Session() as sess:
        path = 'models/crack_capcha.model-' + str(step)
        saver.restore(sess, path)

        predict = tf.argmax(tf.reshape(
            output, [-1, MAX_CAPTCHA, CHAR_SET_LEN]), 2)
        text_list = sess.run(predict, feed_dict={
                             X: [captcha_image], keep_prob: 1})
        text = text_list[0].tolist()
        return text

path=sys.path[0]

if __name__ == '__main__':
    # 训练和测试开关

    train = 0
    if train == 0:

        train_continue(int(ini_read(path+"/step.ini", "value", "step")),
                       float(ini_read(path+"/step.ini", "value", "acc")))

    if train == 1:
        text, image = gen_captcha_text_and_image()

        f = plt.figure()
        ax = f.add_subplot(111)
        ax.text(0.1, 0.9, text, ha='center',
                va='center', transform=ax.transAxes)
        plt.imshow(image)
        plt.show()

        image = convert2gray(image)
        image = image.flatten() / 255

        predict_text = crack_captcha(image, 36300)
        print("正确: {}  预测: {}".format(
            text, [char_set[char] for i, char in enumerate(predict_text)]))
